import numpy as np
from collections import defaultdict

class SimpleBayesClassifier:
    def __init__(self):
        self.class_probs = defaultdict(float)
        self.feature_probs = defaultdict(lambda: defaultdict(float))
    
    def fit(self, X, y):
        """训练分类器"""
        # 计算每个类别的先验概率
        total_samples = len(y)
        for class_label in set(y):
            class_count = sum(1 for label in y if label == class_label)
            self.class_probs[class_label] = class_count / total_samples
        
        # 计算每个特征在每个类别下的条件概率
        for class_label in set(y):
            class_indices = [i for i, label in enumerate(y) if label == class_label]
            class_samples = X[class_indices]
            
            for feature_idx in range(X.shape[1]):
                feature_values = class_samples[:, feature_idx]
                # 使用拉普拉斯平滑避免零概率
                unique_values, counts = np.unique(feature_values, return_counts=True)
                total_class_samples = len(class_samples)
                
                for value, count in zip(unique_values, counts):
                    prob = (count + 1) / (total_class_samples + len(unique_values))
                    self.feature_probs[(feature_idx, value)][class_label] = prob
    
    def predict(self, X):
        """预测新样本"""
        predictions = []
        for sample in X:
            best_class = None
            best_score = -float('inf')
            
            for class_label in self.class_probs:
                # 计算后验概率（取对数避免数值下溢）
                score = np.log(self.class_probs[class_label])  # 先验概率
                
                for feature_idx, value in enumerate(sample):
                    feature_prob = self.feature_probs.get((feature_idx, value), {}).get(class_label, 1e-6)
                    score += np.log(feature_prob)  # 条件概率
                
                if score > best_score:
                    best_score = score
                    best_class = class_label
            
            predictions.append(best_class)
        
        return predictions

# 测试分类器
if __name__ == "__main__":
    # 简单示例数据
    X_train = np.array([
        [1, 0],  # 样本1
        [0, 1],  # 样本2
        [1, 1],  # 样本3
        [0, 0]   # 样本4
    ])
    y_train = np.array([0, 1, 0, 1])  # 类别标签
    
    classifier = SimpleBayesClassifier()
    classifier.fit(X_train, y_train)
    
    X_test = np.array([[1, 0], [0, 0]])
    predictions = classifier.predict(X_test)
    print("预测结果:", predictions)